#include "threadpool.h"

ThreadPool *threadpool_create(const uint32_t workers_count, const uint32_t tasks_capacity) {
    ThreadPool *pool = (ThreadPool *)malloc(sizeof(ThreadPool));
    if (pool == NULL) {
        perror(" - ERROR: thread pool allocation failed");
        exit(EXIT_FAILURE);
    }
    if(workers_count <= 0 || tasks_capacity <= 0) {
        perror(" - ERROR: workers count or tasks capacity is invalid");
        exit(EXIT_FAILURE);
    }

    pool->running = true;
    pool->worker_count = workers_count;
    pool->workers = malloc(workers_count * sizeof(pthread_t));
    if(pool->workers == NULL) {
        perror(" - ERROR: thread workers allocation failed");
        exit(EXIT_FAILURE);
    }
    pool->tasks = queue_create(tasks_capacity);
    pthread_mutex_init(&pool->mutex, NULL);
    pthread_cond_init(&pool->tasks_not_empty, NULL);
    pthread_cond_init(&pool->tasks_not_full, NULL);

    for(uint32_t i = 0; i < pool->worker_count; i++) {
        // 创建工作线程
        if (pthread_create(&pool->workers[i], NULL, threadpool_worker_thread, pool) != 0) {
            perror("pthread_create failed");
            pool->worker_count = i + 1;    // 创建工作线程失败时保存工作线程个数，以便正确回收资源
            threadpool_destroy(pool);
            return NULL;
        }
    }

    return pool;
}

void threadpool_destroy(ThreadPool *pool) {
    if (pool == NULL) {
        return;
    }

    pthread_mutex_lock(&pool->mutex);
    pool->running = false;
    pthread_mutex_unlock(&pool->mutex);

    // 唤醒所有工作线程
    pthread_cond_broadcast(&pool->tasks_not_empty);

    for (uint32_t i = 0; i < pool->worker_count; ++i) {
        // 阻塞等待子线程（工作线程）结束
        pthread_join(pool->workers[i], NULL);
    }

    pool->worker_count = 0;
    free(pool->workers);
    queue_destroy(pool->tasks);
    pthread_mutex_destroy(&pool->mutex);
    pthread_cond_destroy(&pool->tasks_not_empty);
    pthread_cond_destroy(&pool->tasks_not_full);
    free(pool);
}

void *threadpool_worker_thread(void *arg) {
    assert(arg != NULL);
    ThreadPool *pool = (ThreadPool *)arg;

    // 工作线程一直运行
    while(1) {
        Task *task;

        pthread_mutex_lock(&pool->mutex);
        // 线程池运行中且任务队列为空时阻塞等待
        while(pool->running && queue_empty(pool->tasks)) {
            pthread_cond_wait(&pool->tasks_not_empty, &pool->mutex);
        }
        // 线程池停止且任务队列为空时退出工作线程
        if(!pool->running && queue_empty(pool->tasks)) {
            pthread_mutex_unlock(&pool->mutex);
            break;
        }
        task = queue_pop(pool->tasks);    // 从工作队列取出任务
        pthread_mutex_unlock(&pool->mutex);

        //! 注意：这里不需要加锁，避免有锁执行耗时间任务
        task->ret = task->func(task->arg);    // 调用目标任务函数

        pthread_mutex_lock(&task->mutex);
        task->ready = true;      // 标记任务已完成
        pthread_mutex_unlock(&task->mutex);
    }
    return NULL;
}

void threadpool_push(ThreadPool *pool, Task *task) {
    assert(pool != NULL);
    assert(task != NULL);

    pthread_mutex_lock(&pool->mutex);
    // 线程池运行中且任务队列满时阻塞等待
    while(pool->running && queue_full(pool->tasks)) {
        pthread_cond_wait(&pool->tasks_not_full, &pool->mutex);
    }
    if(!pool->running) {
        pthread_mutex_unlock(&pool->mutex);
        return;
    }
    queue_push(pool->tasks, task);
    pthread_mutex_unlock(&pool->mutex);

    // 唤醒一个等待资源的工作线程
    pthread_cond_signal(&pool->tasks_not_empty);
}
